#include <jni.h>
#include <string>
#include "common.h"

#ifdef __cplusplus
extern "C"
{
#endif
    #include <libavformat/avformat.h>
    #include <libavutil/imgutils.h>
    #include <libavutil/pixfmt.h>

#ifdef __cplusplus
};
#endif

FILE * outFp;
AVPacket avpkt;
AVCodecContext *avctx;
const AVCodec *codec;
AVFrame *frame;
uint8_t *video_dst_data[4] = {NULL};
int      video_dst_linesize[4];
int video_dst_bufsize;

void ffmpeg_log_callback(void *ptr, int level, const char *fmt, va_list vl)
{
    va_list vl2;
    char line[1024];
    static int print_prefix = 1;


    va_copy(vl2, vl);
    // av_log_default_callback(ptr, level, fmt, vl);
    av_log_format_line(ptr, level, fmt, vl2, line, sizeof(line), &print_prefix);
    va_end(vl2);

    LOGD("%s", line);

}


extern "C" JNIEXPORT void JNICALL
Java_com_example_king_ffmpegexample_NativeH264Android_init(
        JNIEnv* env,
        jobject /* this */,
        jstring outputFilePath) {

    av_register_all();
    av_log_set_callback(ffmpeg_log_callback);

    const char * video_dst_filename = env->GetStringUTFChars(outputFilePath, NULL);
    if (!video_dst_filename) {
        return;
    }

    outFp = fopen(video_dst_filename, "wb");
    if (!outFp) {
        LOGE("Could not open %s\n", video_dst_filename);
        env->ReleaseStringUTFChars(outputFilePath, video_dst_filename);
        return;
    }

    av_init_packet(&avpkt);

    codec = avcodec_find_decoder(AV_CODEC_ID_H264);
    if (!codec) {
        LOGE("Codec not found\n");
        env->ReleaseStringUTFChars(outputFilePath, video_dst_filename);
        return;
    }
    avctx = avcodec_alloc_context3(codec);
    if (!avctx) {
        LOGE("Could not allocate video codec context\n");
        env->ReleaseStringUTFChars(outputFilePath, video_dst_filename);
        return;
    }
    if (codec->capabilities & AV_CODEC_CAP_TRUNCATED)
        avctx->flags |= AV_CODEC_FLAG_TRUNCATED; // we do not send complete frames

    avctx->width = 240;
    avctx->height = 320;
    avctx->time_base.num = 1;
    avctx->frame_number = 1;
    avctx->codec_type = AVMEDIA_TYPE_VIDEO;
    avctx->bit_rate = 0;

    /* open it */
    if (avcodec_open2(avctx, codec, NULL) < 0) {
        LOGE("Could not open codec\n");
        exit(1);
    }

    frame = av_frame_alloc();
    if (!frame) {
        LOGE("Could not allocate video frame\n");
        env->ReleaseStringUTFChars(outputFilePath, video_dst_filename);
        return;
    }

    int ret = av_image_alloc(video_dst_data, video_dst_linesize,
                         avctx->width, avctx->height, AV_PIX_FMT_YUV420P, 1);

    if (ret < 0) {
        LOGE("Could not allocate raw video buffer\n");
        env->ReleaseStringUTFChars(outputFilePath, video_dst_filename);
        return;
    }
    video_dst_bufsize = ret;


    env->ReleaseStringUTFChars(outputFilePath, video_dst_filename);
}

int saved_flag = -1;

extern "C" JNIEXPORT jint JNICALL
Java_com_example_king_ffmpegexample_NativeH264Android_decodeH264ToYuv(
        JNIEnv* env,
        jobject /* this */,
        jbyteArray in,
        jint len) {

    if (len > 0) {
        jbyte *buf = env->GetByteArrayElements(in, NULL);
        if (buf) {
            avpkt.size = len;
            avpkt.data = (uint8_t *) buf;
            int got_frame = 0;

            int decodedLen = avcodec_decode_video2(avctx, frame, &got_frame, &avpkt);
            if (decodedLen < 0) {
                LOGE("Error while decoding frame");
            }
            if (got_frame) {
                av_image_copy(video_dst_data, video_dst_linesize,
                              (const uint8_t **) (frame->data), frame->linesize,
                              AV_PIX_FMT_YUV420P, avctx->width, avctx->height);

                if (saved_flag < 0) {
                    LOGD("save yuv ---------------------");
                    saved_flag = 1;
                    fwrite(video_dst_data[0], 1, video_dst_bufsize, outFp);
                }
            }

            env->ReleaseByteArrayElements(in, buf, 0);
        }
    }
    return 0;
}

extern "C" JNIEXPORT void JNICALL
        Java_com_example_king_ffmpegexample_NativeH264Android_release(
                JNIEnv* env,
                jobject /* this */) {
    avcodec_free_context(&avctx);
    av_frame_free(&frame);
    av_free(video_dst_data[0]);
    if (outFp) {
        fclose(outFp);
    }

}

